#ifndef FLOW_SOLVER_H
#define FLOW_SOLVER_H

#include <cassert>
#include <cmath>
#include <fstream>
#include <map>
#include <string>
#include <vector>
#include "GridOperator.cuh"
#include "GridPainter.cuh"
#include "GridSolver.cuh"

#ifndef CUDA_SAFE_FREE
#define CUDA_SAFE_FREE(ptr) { if(ptr) { cutilSafeCall(cudaFree(ptr)); ptr = NULL; } }
#endif

template<typename T>
class Field
{
protected:
	int nxy, nxyz;

public:
	int nx, ny, nz;
	T dx;
	FieldType type;
	std::string name;
	std::vector<T> values;
	T *ptr;	// CUDA pointer

public:
	Field()
	{
		nx = ny = nz = 0;
		dx = T(0);
		type = CellCenter;
		name = "Unnamed Field";
		values.clear();
		ptr = NULL;
	}

	~Field()
	{
		CUDA_SAFE_FREE(ptr);
	}

public:
	void load(const std::string &path)
	{
		std::ifstream file(path.c_str(), std::ios::in | std::ios::binary);
		if(!file.is_open())
			return;

		int len = name.length();
		file.read((char *)&len, sizeof(len));
		name.resize(len);
		file.read(&name.at(0), name.size() * sizeof(char));

		file.read((char *)&nx, sizeof(nx));
		file.read((char *)&ny, sizeof(ny));
		file.read((char *)&nz, sizeof(nz));
		file.read((char *)&dx, sizeof(dx));
		file.read((char *)&type, sizeof(type));
		resize(nx, ny, nz, dx, type);

		file.read((char *)&values.front(), values.size() * sizeof(T));
		file.close();
		syncHostToDevice();
	}

	void loadRAW16I(const std::string &path, int width, int height, int depth)
	{
		std::ifstream file(path.c_str(), std::ios::in | std::ios::binary);
		if(!file.is_open())
			return;

		name = path;
		type = CellCenter;
		nx = width;
		ny = height;
		nz = depth;
		dx = 1.0f;
		resize(nx, ny, nz, dx, type);
		std::vector<unsigned short> buffer(values.size());
		file.read((char *)&buffer.front(), buffer.size() * sizeof(unsigned short));
		file.close();
		for(size_t i = 0; i < values.size(); ++i)
			values[i] = T(buffer[i]) / T(4096);
		syncHostToDevice();
	}

	void save(const std::string &path)
	{
		syncDeviceToHost();
		std::ofstream file(path.c_str(), std::ios::out | std::ios::binary);
		if(!file.is_open())
			return;

		int len = name.length();
		file.write((const char *)&len, sizeof(len));
		file.write(name.c_str(), name.size() * sizeof(char));

		file.write((const char *)&nx, sizeof(nx));
		file.write((const char *)&ny, sizeof(ny));
		file.write((const char *)&nz, sizeof(nz));
		file.write((const char *)&dx, sizeof(dx));
		file.write((const char *)&type, sizeof(type));

		file.write((const char *)&values.front(), values.size() * sizeof(T));
		file.close();
	}

	void access(int iMin, int jMin, int kMin, int iMax, int jMax, int kMax, void accessor(Field<T> &, int i, int j, int k))
	{
		assert(iMin >= 0 && iMax < nx && jMin >= 0 && jMax < ny && kMin >= 0 && kMax < nz);
		for(int k = kMin; k <= kMax; ++k)
			for(int j = jMin; j <= jMax; ++j)
				for(int i = iMin; i <= iMax; ++i)
					accessor(*this, i, j, k);
	}

	void resize(int nx, int ny, int nz, T dx, const FieldType type)
	{
		this->nx = nx;
		this->ny = ny;
		this->nz = nz;
		this->dx = dx;
		this->type = type;
		nxy = nx * ny;
		nxyz = nxy * nz;

		CUDA_SAFE_FREE(ptr);
		if(nxyz <= 0)
		{
			values.clear();
			return;
		}

		values.assign(nxyz, T(0));
		cutilSafeCall(cudaMalloc((void **)&ptr, nxyz * sizeof(T)));
		cutilSafeCall(cudaMemset(ptr, 0, nxyz * sizeof(T)));
		std::cout << name << ": " << nx << " x " << ny << " x " << nz << std::endl;
	}

	void copyFromDevice(T *src)
	{
		if(nxyz > 0)
			cutilSafeCall(cudaMemcpy(ptr, src, nxyz * sizeof(T), cudaMemcpyDeviceToDevice));
	}

	void syncHostToDevice()
	{
		if(nxyz > 0)
			cutilSafeCall(cudaMemcpy(ptr, &values[0], nxyz * sizeof(T), cudaMemcpyHostToDevice));
	}

	void syncDeviceToHost()
	{
		if(nxyz > 0)
			cutilSafeCall(cudaMemcpy(&values[0], ptr, nxyz * sizeof(T), cudaMemcpyDeviceToHost));
	}

	inline T &operator[](int idx)
	{
		assert(idx >= 0 && idx < nxyz);
		return values[idx];
	}

	inline T &operator()(int i, int j, int k)
	{
		assert(i >= 0 && i < nx && j >= 0 && j < ny && k >= 0 && k < nz);
		return values[i + j * nx + k * nxy];
	}

	inline const T evaluate(T x, T y, T z) const
	{
		T rdx = T(1) / dx;
		x *= rdx;
		y *= rdx;
		z *= rdx;

		switch(type)
		{
		case CellCenter:
			x -= T(0.5);
			y -= T(0.5);
			z -= T(0.5);
			break;
		case FaceCenterX:
			y -= T(0.5);
			z -= T(0.5);
			break;
		case FaceCenterY:
			x -= T(0.5);
			z -= T(0.5);
			break;
		case FaceCenterZ:
			x -= T(0.5);
			y -= T(0.5);
			break;
		}

		if(x < T(0)) x = T(0); if(x > T(nx - 1)) x = T(nx - 1);
		if(y < T(0)) y = T(0); if(y > T(ny - 1)) y = T(ny - 1);
		if(z < T(0)) z = T(0); if(z > T(nz - 1)) z = T(nz - 1);
		int xi(x), yi(y), zi(z);
		T xf(x - xi), yf(y - yi), zf(z - zi);
		while(xi >= nx - 1) { --xi; xf += T(1); }
		while(yi >= ny - 1) { --yi; yf += T(1); }
		while(zi >= nz - 1) { --zi; zf += T(1); }
		int idx = xi + yi * nx + zi * nxy;
		T vx1 = values[idx] * (T(1) - xf) + values[idx + 1] * xf;
		T vx2 = values[idx + nx] * (T(1) - xf) + values[idx + nx + 1] * xf;
		idx += nxy;
		T vx3 = values[idx] * (T(1) - xf) + values[idx + 1] * xf;
		T vx4 = values[idx + nx] * (T(1) - xf) + values[idx + nx + 1] * xf;
		T vy1 = vx1 * (T(1) - yf) + vx2 * yf;
		T vy2 = vx3 * (T(1) - yf) + vx4 * yf;
		return vy1 * (T(1) - zf) + vy2 * zf;
	}
};

/**
	@brief basic field and particle operators for fluid simulation
*/
template<typename T>
class FlowSolver
{
protected:
	int its;
	std::ofstream log;
	std::map<std::string, Field<T>> fields;
	std::map<std::string, std::vector<T>> attributes;

public:
	int nx, ny, nz;
	T dx;

public:
	FlowSolver() : log("FlowSolver.log")
	{
		its = 32;
		nx = ny = nz = 0;
		dx = T(0);
	}

public:
	void resize(int nx, int ny, int nz, T dx)
	{
		this->nx = nx;
		this->ny = ny;
		this->nz = nz;
		this->dx = dx;

		// resize all fields
		for(std::map<std::string, Field<T>>::iterator it = fields.begin(); it != fields.end(); ++it)
			resizeField(it->second);
	}

	void resizeField(Field<T> &field)
	{
		switch(field.type)
		{
		case CellCenter:
			field.resize(nx, ny, nz, dx, field.type);
			break;
		case FaceCenterX:
			field.resize(nx + 1, ny, nz, dx, field.type);
			break;
		case FaceCenterY:
			field.resize(nx, ny + 1, nz, dx, field.type);
			break;
		case FaceCenterZ:
			field.resize(nx, ny, nz + 1, dx, field.type);
			break;
		}
	}

	Field<T> *addField(const std::string &name, const FieldType type = CellCenter, bool replace = false)
	{
		std::map<std::string, Field<T>>::iterator it = fields.find(name);
		if(it != fields.end() && !replace)
		{
			log << "[Error] addField() - field '" << name << "' already exist!" << std::endl;
			return NULL;
		}

		Field<T> &field = fields[name];
		field.type = type;
		field.name = name;
		resizeField(field);
		return &field;
	}

	void dropField(const std::string &name)
	{
		fields.erase(name);
	}

	void setBoundaryConditionX(Field<T> &f, T bc)
	{
		applyBoundaryX(f.nx, f.ny, f.nz, bc, f.ptr);
	}

	void setBoundaryConditionY(Field<T> &f, T bc)
	{
		applyBoundaryY(f.nx, f.ny, f.nz, bc, f.ptr);
	}

	void setBoundaryConditionZ(Field<T> &f, T bc)
	{
		applyBoundaryZ(f.nx, f.ny, f.nz, bc, f.ptr);
	}

	void evaluateCurl(const Field<T> &u, const Field<T> &v, const Field<T> &w, Field<T> &curlx, Field<T> &curly, Field<T> &curlz)
	{
	}

	void evaluateDivergence(const Field<T> &u, const Field<T> &v, const Field<T> &w, Field<T> &div)
	{
		divergence(nx, ny, nz, dx, u.ptr, v.ptr, w.ptr, div.ptr);
	}

	void evaluateGradient(const Field<T> &f, Field<T> &gradx, Field<T> &grady, Field<T> &gradz)
	{
	}

	void evaluateLaplacian(const Field<T> &f, Field<T> &lapl)
	{
	}

	void solveConvection(T dt, const Field<T> &u, const Field<T> &v, const Field<T> &w, const Field<T> &f, Field<T> &fout, const NumericalScheme scheme = SemiLagrange)
	{
		convect(fout.nx, fout.ny, fout.nz, dx, dt, u.ptr, v.ptr, w.ptr, f.ptr, fout.ptr, fout.type, scheme);
	}

	void solveDiffusion(T dt, T kappa, const Field<T> &f, Field<T> &fout)
	{
		solveHeatEquation(nx, ny, nz, dx, dt, kappa, f.ptr, fout.ptr, its);
	}

	void solveProjection(Field<T> &u, Field<T> &v, Field<T> &w, Field<T> &div, Field<T> &p)
	{
		// During pressure-Poisson projection, we call divergence using dx=-1
		// This trick can eliminate -1/dx^2 from Laplacian operator and 1/dx from gradient operator
		divergence(nx, ny, nz, T(-1), u.ptr, v.ptr, w.ptr, div.ptr);
		solvePoissonEquation(nx, ny, nz, T(1), div.ptr, p.ptr, its);
		project(nx, ny, nz, T(1), u.ptr, v.ptr, w.ptr, p.ptr);
	}
};

template<typename T>
class FireSolver : public FlowSolver<T>
{
public:
	T uBC, vBC, wBC;
	T accTemp, accTurb, ambTemp, maxTemp;
	T burnRate, coolRate;
	Field<T> *u0, *u;	// velocity in x
	Field<T> *v0, *v;	// velocity in y
	Field<T> *w0, *w;	// velocity in z
	Field<T> *f0, *f;	// fuel
	Field<T> *s0, *s;	// soot
	Field<T> *t0, *t;	// temperature
	Field<T> *d, *p;	// divergence and pressure
	Field<T> *fx;		// velocity/turbulence in x at cell center
	Field<T> *fy;		// velocity/turbulence in y at cell center
	Field<T> *fz;		// velocity/turbulence in z at cell center
	Field<T> *cx;		// curl in x at cell center
	Field<T> *cy;		// curl in y at cell center
	Field<T> *cz;		// curl in z at cell center
	Field<T> *cn;		// curl norm at cell center

public:
	FireSolver()
	{
		uBC = T(0);
		vBC = T(0);
		wBC = T(0);
		accTemp = T(0);
		accTurb = T(0);
		ambTemp = T(0);
		maxTemp = T(0);
		burnRate = T(0);
		coolRate = T(0);
		u0 = addField("u0", FaceCenterX);
		v0 = addField("v0", FaceCenterY);
		w0 = addField("w0", FaceCenterZ);
		f0 = addField("f0", CellCenter);
		s0 = addField("s0", CellCenter);
		t0 = addField("t0", CellCenter);
		u = addField("u", FaceCenterX);
		v = addField("v", FaceCenterY);
		w = addField("w", FaceCenterZ);
		f = addField("f", CellCenter);
		s = addField("s", CellCenter);
		t = addField("t", CellCenter);
		d = addField("d", CellCenter);
		p = addField("p", CellCenter);
		fx = addField("fx", CellCenter);
		fy = addField("fy", CellCenter);
		fz = addField("fz", CellCenter);
		cx = addField("cx", CellCenter);
		cy = addField("cy", CellCenter);
		cz = addField("cz", CellCenter);
		cn = addField("cn", CellCenter);
	}

public:
	void drawCylinder(T cx, T cy, T cz, T r, T h, T v, Field<T> *field)
	{
		cylinder(nx, ny, nz, dx, cx, cy, cz, r, h, v, field->ptr);
	}

	void drawSphere(T cx, T cy, T cz, T r, T v, Field<T> *field)
	{
		sphere(nx, ny, nz, dx, cx, cy, cz, r, v, field->ptr);
	}

	void drawRectangle(T cx, T cy, T cz, T w, T h, T d, T v, Field<T> *field)
	{
		rectangle(nx, ny, nz, dx, cx, cy, cz, w, h, d, v, field->ptr);
	}

	void step(T dt)
	{
		// swap fields
		Field<T> *temp;
		temp = u0; u0 = u; u = temp;
		temp = v0; v0 = v; v = temp;
		temp = w0; w0 = w; w = temp;
		temp = f0; f0 = f; f = temp;
		temp = s0; s0 = s; s = temp;
		temp = t0; t0 = t; t = temp;

		// solve convection for all fields
		solveConvection(dt, *u0, *v0, *w0, *u0, *u);
		solveConvection(dt, *u0, *v0, *w0, *v0, *v);
		solveConvection(dt, *u0, *v0, *w0, *w0, *w);
		solveConvection(dt, *u0, *v0, *w0, *f0, *f);
		solveConvection(dt, *u0, *v0, *w0, *s0, *s);
		solveConvection(dt, *u0, *v0, *w0, *t0, *t);

		// solve diffusion
		//temp = t0; t0 = t; t = temp;
		//solveDiffusion(dt, T(0.1), *t0, *t);

		// apply buoyancy and reaction
		gravity(v->nx, v->ny, v->nz, dx, dt, accTemp, ambTemp, t->ptr, v->ptr);
		t->syncDeviceToHost();
		maxTemp = T(0);
		for(size_t i = 0; i < t->values.size(); ++i)
			maxTemp = std::max(maxTemp, t->values[i]);
		reaction(nx, ny, nz, dx, dt, burnRate, coolRate, ambTemp, maxTemp, f->ptr, s->ptr, t->ptr);

		// apply vorticity confinement
		resample(nx, ny, nz, dx, u0->ptr, v0->ptr, w0->ptr, fx->ptr, fy->ptr, fz->ptr);
		curl(nx, ny, nz, T(1), fx->ptr, fy->ptr, fz->ptr, cx->ptr, cy->ptr, cz->ptr);
		l2norm(nx, ny, nz, dx, cx->ptr, cy->ptr, cz->ptr, cn->ptr);
		turbulence(nx, ny, nz, T(1), accTurb, cx->ptr, cy->ptr, cz->ptr, cn->ptr, fx->ptr, fy->ptr, fz->ptr);
		accelerate(nx, ny, nz, dx, dt, fx->ptr, fy->ptr, fz->ptr, u->ptr, v->ptr, w->ptr);

		// apply boundary condition for velocity field
		setBoundaryConditionX(*u, uBC);
		setBoundaryConditionY(*v, vBC);
		setBoundaryConditionZ(*w, wBC);

		// apply divergence free constraint
		//solveProjection(*u, *v, *w, *d, *p);
		divergence(nx, ny, nz, T(-1), u->ptr, v->ptr, w->ptr, d->ptr);
		//add(nx, ny, nz, T(1), d->ptr, T(10) * dx, t->ptr, d->ptr);
		solvePoissonEquation(nx, ny, nz, T(1), d->ptr, p->ptr, its);
		project(nx, ny, nz, T(1), u->ptr, v->ptr, w->ptr, p->ptr);
	}
};

#endif
